# -*- coding: utf-8 -*-
# 标准库导入
import os
import json
import logging
import traceback
import base64
from datetime import datetime

# 第三方库导入
import websocket
import ssl
from flask import Flask, request, render_template, jsonify, send_from_directory, abort, url_for, redirect
from werkzeug.utils import secure_filename
from flask_login import LoginManager, login_user, logout_user, login_required, current_user

# 本地模块导入
from config import *
from chat_module import chat_with_deepseek
from vector_db import VectorDB
from quiz_module import QuizModule
from image_recognition import ImageRecognition
from history_module import HistoryModule
from cache_module import CacheModule
from garbage_labels import GarbageLabels
from search_module import SearchModule
from file_utils import FileUtils
from url_utils import URLUtils
from db_config import db, SQLALCHEMY_DATABASE_URI, SQLALCHEMY_TRACK_MODIFICATIONS
from user_module import User, HistoryRecord
from ImageUnderstanding import main as image_understanding_main, continue_dialog as image_understanding_continue
from spark_config import SPARK_CONFIG, WS_URL

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化全局变量
_is_initialized = False
answer = ""  # 用于存储星火API的回答

# 配置文件上传
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

# 创建 Flask 应用
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传文件大小为16MB
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 配置
MODEL_PATH = 'checkpoints/model_best_checkpoint_resnet50.pth.tar'
DATASET_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), '垃圾图片库')
CACHE_DIR = 'feature_cache'

# 创建实例
vector_db = VectorDB()
quiz_module = QuizModule()
image_recognition = None  # 将在初始化时创建
history_module = HistoryModule(UPLOAD_FOLDER)
cache_module = CacheModule(CACHE_DIR, DATASET_PATH)
garbage_labels = GarbageLabels()
search_module = SearchModule(DATASET_PATH)
file_utils = FileUtils(UPLOAD_FOLDER, ALLOWED_EXTENSIONS)
url_utils = URLUtils(DATASET_PATH, UPLOAD_FOLDER)  # 添加 upload_folder 参数

# 初始化Flask-Login
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'welcome'  # 修改为welcome页面
login_manager.login_message = '请先登录后再访问此页面'
login_manager.login_message_category = 'info'

# 配置数据库
app.config['SQLALCHEMY_DATABASE_URI'] = SQLALCHEMY_DATABASE_URI
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = SQLALCHEMY_TRACK_MODIFICATIONS
app.config['SECRET_KEY'] = 'your-secret-key-here'  # 用于session加密
db.init_app(app)

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

# 添加聊天接口
@app.route('/chat', methods=['POST'])
def chat():
    """处理聊天请求，支持流式响应和思考过程"""
    try:
        data = request.get_json()
        user_message = data.get('message', '')
        
        if not user_message:
            return jsonify({'error': '请输入问题'}), 400
            
        # 使用生成器返回流式响应
        return chat_with_deepseek(user_message)
        
    except Exception as e:
        logger.error(f'处理聊天请求时出错: {str(e)}')
        logger.error(f'错误堆栈: {traceback.format_exc()}')
        return jsonify({'error': '服务器处理请求时出错'}), 500

# 创建必要的目录
def ensure_directories():
    """确保必要的目录存在"""
    try:
        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
        os.makedirs('user_history', exist_ok=True)
        logger.info("必要的目录已创建")
    except Exception as e:
        logger.error(f"创建目录失败: {str(e)}")

# 确保在应用启动时创建必要的目录
ensure_directories()

# 全局变量
image_recognition = None  # 将在初始化时创建

@app.route('/dataset/<path:filename>')
def serve_dataset(filename):
    """提供数据集图片访问服务"""
    return url_utils.serve_dataset_file(filename)

@app.route('/')
@login_required
def home():
    return render_template('index.html')

@app.route('/welcome')
def welcome():
    if (current_user.is_authenticated):
        return redirect(url_for('home'))
    return render_template('welcome.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if (request.method == 'POST'):
        try:
            username = request.form.get('username')
            password = request.form.get('password')
            
            if (not username or not password):
                return jsonify({'success': False, 'error': '用户名和密码不能为空'})
            
            if (User.get_by_username(username)):
                return jsonify({'success': False, 'error': '用户名已存在'})
            
            user = User(username=username)
            user.set_password(password)
            
            db.session.add(user)
            db.session.commit()
            
            return jsonify({'success': True, 'message': '注册成功'})
        except Exception as e:
            db.session.rollback()
            return jsonify({'success': False, 'error': str(e)})
    return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if (request.method == 'POST'):
        username = request.form.get('username')
        password = request.form.get('password')
        
        user = User.get_by_username(username)
        if (user and user.check_password(password)):
            login_user(user)
            user.last_login = datetime.utcnow()
            db.session.commit()
            return jsonify({'success': True})
        return jsonify({'success': False, 'error': '用户名或密码错误'})
    return render_template('login.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('home'))

@app.route('/predict', methods=['POST'])
@login_required
def predict():
    if ('file' not in request.files):
        return jsonify({'error': '没有文件上传'})
    
    file = request.files['file']
    if (file.filename == ''):
        return jsonify({'error': '没有选择文件'})
    
    if (not file_utils.allowed_file(file.filename)):
        return jsonify({'error': '不支持的文件类型'})
    
    try:
        # 保存上传的文件
        filepath, filename = file_utils.save_uploaded_file(file)
        logger.info(f"文件已保存: {filepath}")
        
        # 进行预测
        result = image_recognition.predict(filepath)
        if (not result['success']):
            return jsonify({'error': '预测失败'})
        
        # 获取预测信息
        prediction_info = garbage_labels.get_prediction_info(result['predicted_idx'])
        if (not prediction_info):
            return jsonify({'error': '获取预测信息失败'})
        
        # 构建图片URL
        image_url = url_for('serve_uploaded_file', filename=filename, _external=True)
        
        # 保存预测历史到数据库
        history_record = HistoryRecord(
            user_id=current_user.id,
            image_path=filepath,
            prediction_result={
                'main_class': prediction_info['main_class'],
                'detailed_class': prediction_info['detailed_class'],
                'confidence': result['confidence']
            }
        )
        db.session.add(history_record)
        db.session.commit()
        
        # 构建元数据并存储到向量数据库
        metadata = {
            'image_url': image_url,
            'main_class': prediction_info['main_class'],
            'detailed_class': prediction_info['detailed_class'],
            'timestamp': datetime.now().isoformat(),
            'confidence': result['confidence'],
            'user_id': current_user.id
        }
        vector_db.store_history_vector(filepath, result['features'], metadata)
        
        # 查找历史记录中的相似图片，并传入当前图片URL排除自身
        similar_history = vector_db.find_similar_in_history(
            result['features'], 
            current_image_url=image_url
        )
        
        # 保存预测结果和历史记录
        response = history_module.save_prediction_with_similar(
            filepath, filename, result, prediction_info, similar_history, image_url
        )
        
        return jsonify(response)
        
    except Exception as e:
        logger.error(f'预测出错: {str(e)}')
        logger.error(f'错误堆栈: {traceback.format_exc()}')
        return jsonify({'error': f'预测出错: {str(e)}'})

@app.route('/uploads/<filename>')
def serve_uploaded_file(filename):
    """处理上传文件的访问"""
    return url_utils.serve_uploaded_file(filename)

@app.route('/text_search', methods=['POST'])
def text_search():
    """文字搜索接口"""
    try:
        data = request.get_json()
        if (not data or 'query' not in data):
            return jsonify({'error': '没有提供搜索词'})
        
        query = data['query'].strip()
        if (not query):
            return jsonify({'error': '搜索词不能为空'})
        
        logger.info(f"\n收到文字搜索请求：{query}")
        results = search_module.search_by_text(query, cache_module, garbage_labels)
        logger.info(f"找到 {len(results)} 个结果")
        
        return jsonify({
            'success': True,
            'results': results
        })
        
    except Exception as e:
        logger.error(f'文字搜索出错: {str(e)}')
        return jsonify({'error': f'搜索出错: {str(e)}'})

@app.route('/search_similar', methods=['POST'])
def search_similar():
    try:
        if ('file' not in request.files):
            return jsonify({'success': False, 'error': '没有上传图片'})
            
        file = request.files['file']
        if (file.filename == ''):
            return jsonify({'success': False, 'error': '没有选择图片'})
            
        if (not file_utils.allowed_file(file.filename)):
            return jsonify({'success': False, 'error': '不支持的文件类型'})
            
        # 保存上传的临时图片
        filepath, _ = file_utils.save_uploaded_file(file, temp=True)
        
        try:
            # 使用image_recognition进行预测和特征提取
            result = image_recognition.predict(filepath)
            if (not result['success']):
                return jsonify({'success': False, 'error': '预测失败'})
                
            # 获取预测的类别信息
            prediction_info = garbage_labels.get_prediction_info(result['predicted_idx'])
            if (not prediction_info):
                return jsonify({'success': False, 'error': '获取预测信息失败'})
                
            # 使用 search_module 查找相似图片
            similar_images = search_module.find_similar_images(
                result['features'], 
                cache_module, 
                garbage_labels,
                vector_db=vector_db
            )
            
            # 处理相似图片的路径，转换为URL
            processed_similar_images = url_utils.process_similar_images(similar_images)
            
            # 返回完整的预测信息
            response_data = {
                'success': True,
                'main_class': prediction_info['main_class'],
                'detailed_class': prediction_info['detailed_class'],
                'definition': prediction_info.get('definition', ''),
                'rules': prediction_info.get('rules', ''),
                'confidence': result['confidence'],
                'similar_images': processed_similar_images
            }
            
            return jsonify(response_data)
            
        except Exception as e:
            logger.error(f"处理图片时出错: {str(e)}")
            return jsonify({'success': False, 'error': f'处理图片时出错: {str(e)}'})
            
        finally:
            # 清理临时文件
            file_utils.remove_file(filepath)
                    
    except Exception as e:
        logger.error(f"搜索相似图片时出错: {str(e)}")
        return jsonify({'success': False, 'error': f'搜索相似图片时出错: {str(e)}'})

@app.route('/get_questions', methods=['GET'])
def get_questions():
    """获取问答题目"""
    result = quiz_module.get_questions()
    return jsonify(result)

@app.route('/submit_quiz', methods=['POST'])
def submit_quiz():
    try:
        data = request.get_json()
        answer = data.get('answer', '')
        correct_answer = data.get('correct_answer', '')
        
        result = quiz_module.check_answer(answer, correct_answer)
        return jsonify(result)
    except Exception as e:
        logger.error(f'提交答案出错: {str(e)}')
        return jsonify({"error": f'提交答案出错: {str(e)}'})

@app.after_request
def add_security_headers(response):
    response.headers['X-Content-Type-Options'] = 'nosniff'
    response.headers['X-Frame-Options'] = 'SAMEORIGIN'
    response.headers['X-XSS-Protection'] = '1; mode=block'
    return response

@app.route('/history', methods=['GET'])
def get_history():
    try:
        history_entries = history_module.get_history(garbage_labels, url_for)
        if (isinstance(history_entries, list)):
            return jsonify(history_entries)
        else:
            return jsonify(history_entries), 500
    except Exception as e:
        logger.error(f"获取历史记录时出错: {str(e)}")
        return jsonify({'error': '获取历史记录失败'}), 500

@app.route('/delete_history/<filename>', methods=['DELETE'])
def delete_history(filename):
    try:
        result = history_module.delete_history(filename)
        if ('error' in result):
            return jsonify(result), result.get('status_code', 500)
        return jsonify(result)
    except Exception as e:
        logger.error(f"删除历史记录时出错: {str(e)}")
        return jsonify({'error': '删除失败'}), 500

@app.route('/get_search_suggestions')
def get_search_suggestions():
    """获取搜索提示词"""
    try:
        return jsonify({
            'success': True,
            'suggestions': garbage_labels.get_search_suggestions()
        })
    except Exception as e:
        logger.error(f"获取搜索提示词失败: {str(e)}")
        return jsonify({
            'error': '获取搜索提示词失败'
        }), 500

def preload_dataset_features():
    """预加载数据集特征"""
    if (not os.path.exists(DATASET_PATH)):
        logger.error(f'错误: 找不到数据集目录 {DATASET_PATH}')
        return
    
    # 获取所有图片文件
    image_files = []
    for root, _, files in os.walk(DATASET_PATH):
        for file in files:
            if (file.lower().endswith(('.png', '.jpg', '.jpeg'))):
                image_files.append(os.path.join(root, file))
    
    if (not image_files):
        logger.warning('警告: 数据集目录中没有找到图片文件')
        return
    
    total_files = len(image_files)
    logger.info(f'找到 {total_files} 个图片文件')
    
    # 设置向量数据库的总向量数，用于显示进度
    vector_db.set_total_vectors(total_files)
    
    cached_count = 0
    new_count = 0
    error_count = 0
    processed_count = 0
    
    # 创建进度条，每10%显示一次
    progress_interval = max(1, total_files // 10)
    
    for image_path in image_files:
        try:
            processed_count += 1
            if (processed_count % progress_interval == 0):
                progress = (processed_count / total_files) * 100
                logger.info(f'进度: {progress:.0f}% ({processed_count}/{total_files})')
            
            # 检查缓存
            features = cache_module.load_feature_cache(image_path)
            if (features is not None):
                cached_count += 1
            else:
                # 使用 image_recognition 提取特征
                features = image_recognition.extract_features(image_path)
                if (features is None):
                    error_count += 1
                    continue
                
                # 保存缓存
                cache_module.save_feature_cache(image_path, features)
                new_count += 1
            
            # 不管是从缓存加载还是新提取的特征，都存储到向量数据库
            class_info = garbage_labels.get_class_from_path(image_path, DATASET_PATH)
            metadata = {
                "class": class_info["class_name"] if class_info else "unknown",
                "main_category": class_info["main_category"] if class_info else "unknown",
                "path": image_path
            }
            vector_db.store_vector(image_path, features, metadata)
            
        except Exception as e:
            logger.error(f'处理图片出错 {image_path}: {str(e)}')
            error_count += 1
            continue
    
    # 总是显示加载统计信息
    logger.info(f'数据集加载完成:')
    logger.info(f'总图片数: {total_files}')
    logger.info(f'缓存加载: {cached_count}')
    logger.info(f'新处理: {new_count}')
    logger.info(f'错误: {error_count}')
    
    # 获取集合信息，确认向量是否成功存储
    try:
        collection_info = vector_db.qdrant_client.get_collection("image_features")
        logger.info(f'向量数据库存储状态: 共 {collection_info.vectors_count} 个向量')
    except Exception as e:
        logger.error(f'获取向量数据库信息失败: {str(e)}')

def initialize_app():
    """初始化应用程序"""
    global image_recognition, _is_initialized
    
    # 检查是否已经初始化
    if (_is_initialized):
        return True
        
    try:
        # 加载标签映射
        if (not garbage_labels.load_label_mappings()):
            logger.error("加载标签映射失败")
            return False
        
        # 创建并初始化图像识别模块
        image_recognition = ImageRecognition(MODEL_PATH, len(garbage_labels.label_to_category))
        if (not image_recognition.initialize()):
            logger.error("图像识别模块初始化失败")
            return False
            
        # 初始化向量数据库
        vector_db.initialize()
        
        # 加载缓存的特征
        cache_module.load_all_features()
        
        # 预加载数据集特征
        preload_dataset_features()
        
        # 设置初始化标志
        _is_initialized = True
        return True
    except Exception as e:
        logger.error(f"初始化应用程序时出错: {str(e)}")
        return False

@app.route('/image_understanding', methods=['POST'])
@login_required
def image_understanding():
    """处理图片理解请求"""
    try:
        if ('file' not in request.files):
            return jsonify({'error': '没有文件上传'})
        
        file = request.files['file']
        if (file.filename == ''):
            return jsonify({'error': '没有选择文件'})
        
        if (not file_utils.allowed_file(file.filename)):
            return jsonify({'error': '不支持的文件类型'})

        # 直接读取上传的文件数据，不保存到磁盘
        file_data = file.read()
        image_base64 = base64.b64encode(file_data).decode('utf-8')

        # 构建第一次固定的问题
        first_question = "图中都有哪些物品，都属于哪类垃圾，其对应的分类说明和处理投放规则是什么？"
        text = [
            {"role": "user", "content": image_base64, "content_type": "image"},
            {"role": "user", "content": first_question}
        ]

        # 从配置中获取参数
        appid = SPARK_CONFIG['APPID']
        api_key = SPARK_CONFIG['APIKey']
        api_secret = SPARK_CONFIG['APISecret']
        
        # 调用 ImageUnderstanding 模块进行图片理解
        response = image_understanding_main(appid, api_key, api_secret, WS_URL, text)
        
        return jsonify({
            'success': True,
            'answer': response
        })
        
    except Exception as e:
        logger.error(f'图片理解出错: {str(e)}')
        return jsonify({'error': f'图片理解出错: {str(e)}'})

@app.route('/image_understanding_follow_up', methods=['POST'])
@login_required
def image_understanding_follow_up():
    """处理图片理解的后续问题"""
    try:
        data = request.get_json()
        follow_up_question = data.get('message', '')
        
        if (not follow_up_question):
            return jsonify({'error': '请输入问题'})

        # 从配置中获取参数
        appid = SPARK_CONFIG['APPID']
        api_key = SPARK_CONFIG['APIKey']
        api_secret = SPARK_CONFIG['APISecret']
        
        # 调用 ImageUnderstanding 模块继续对话
        response = image_understanding_continue(
            appid, api_key, api_secret, WS_URL, follow_up_question
        )
        
        return jsonify({
            'success': True,
            'answer': response
        })
        
    except Exception as e:
        logger.error(f'处理后续问题出错: {str(e)}')
        return jsonify({'error': f'处理出错: {str(e)}'})

if __name__ == '__main__':
    import argparse
    
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='废弃物数据管理系统服务器')
    parser.add_argument('--debug', action='store_true', 
                      help='是否使用调试模式运行（默认: 否）')
    args = parser.parse_args()
    
    # 初始化应用
    initialize_app()
    
    # 根据命令行参数决定是否使用调试模式
    if (args.debug):
        app.run(debug=True, use_reloader=False)
    else:
        app.run(host='127.0.0.1', port=5000)
